import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (classification_report, confusion_matrix, 
                            accuracy_score, precision_score, recall_score, 
                            f1_score, roc_auc_score, precision_recall_curve)
from sklearn.model_selection import GridSearchCV
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def load_processed_data():
    """加载预处理后的数据"""
    print("=== 加载数据 ===")
    
    try:
        # 加载训练数据
        train_data = pd.read_excel('data_train_balanced.xlsx')
        test_data = pd.read_excel('data_test_validation.xlsx')
        
        print(f"训练集: {len(train_data)}行, {len(train_data.columns)}列")
        print(f"测试集: {len(test_data)}行, {len(test_data.columns)}列")
        
        # 分离特征和标签
        feature_columns = [col for col in train_data.columns if col != '是否异常']
        
        X_train = train_data[feature_columns]
        y_train = train_data['是否异常']
        X_test = test_data[feature_columns]
        y_test = test_data['是否异常']
        
        print(f"特征数量: {len(feature_columns)}")
        print(f"训练集分布 - 正常: {(y_train == 0).sum()}, 异常: {(y_train == 1).sum()}")
        print(f"测试集分布 - 正常: {(y_test == 0).sum()}, 异常: {(y_test == 1).sum()}")
        
        return X_train, X_test, y_train, y_test, feature_columns
        
    except Exception as e:
        print(f"数据加载失败: {e}")
        return None, None, None, None, None

def train_optimized_random_forest(X_train, y_train, random_state=42):
    """训练优化的随机森林模型"""
    print("\n=== 训练优化随机森林模型 ===")
    
    # 使用之前找到的最佳参数
    best_params = {
        'max_depth': None, 
        'min_samples_leaf': 1, 
        'min_samples_split': 2, 
        'n_estimators': 200
    }
    
    rf_optimized = RandomForestClassifier(
        random_state=random_state,
        n_jobs=-1,
        **best_params
    )
    
    rf_optimized.fit(X_train, y_train)
    print("优化随机森林模型训练完成")
    
    return rf_optimized

def find_optimal_threshold_for_recall(model, X_test, y_test, min_recall=0.90):
    """寻找最优阈值以达到目标召回率"""
    print(f"\n=== 寻找最优阈值（目标召回率≥{min_recall:.0%}） ===")
    
    # 获取预测概率
    y_proba = model.predict_proba(X_test)[:, 1]
    
    # 计算精确率-召回率曲线
    precision, recall, thresholds = precision_recall_curve(y_test, y_proba)
    
    # 寻找满足最小召回率要求的阈值
    valid_indices = recall >= min_recall
    
    if not any(valid_indices):
        print(f"警告: 无法达到{min_recall:.0%}的召回率")
        # 选择召回率最高的阈值
        max_recall_idx = np.argmax(recall)
        optimal_threshold = thresholds[max_recall_idx] if max_recall_idx < len(thresholds) else 0.0
        optimal_recall = recall[max_recall_idx]
        optimal_precision = precision[max_recall_idx]
        print(f"最高可达召回率: {optimal_recall:.4f}")
    else:
        # 在满足召回率要求的阈值中，选择精确率最高的
        # 注意：thresholds长度比precision和recall少1，需要调整索引
        valid_indices_adj = valid_indices[:-1]  # 调整索引以匹配thresholds
        
        if not any(valid_indices_adj):
            # 如果调整后没有有效索引，使用最低阈值
            optimal_threshold = 0.0
            max_recall_idx = np.argmax(recall)
            optimal_recall = recall[max_recall_idx]
            optimal_precision = precision[max_recall_idx]
        else:
            valid_precisions = precision[:-1][valid_indices_adj]
            valid_recalls = recall[:-1][valid_indices_adj]
            valid_thresholds = thresholds[valid_indices_adj]
            
            best_idx = np.argmax(valid_precisions)
            optimal_threshold = valid_thresholds[best_idx]
            optimal_recall = valid_recalls[best_idx]
            optimal_precision = valid_precisions[best_idx]
        
        print(f"找到满足条件的阈值: {optimal_threshold:.4f}")
    
    print(f"对应的召回率: {optimal_recall:.4f}")
    print(f"对应的精确率: {optimal_precision:.4f}")
    print(f"对应的F1得分: {2 * optimal_precision * optimal_recall / (optimal_precision + optimal_recall):.4f}")
    
    return optimal_threshold, optimal_recall, optimal_precision

def evaluate_different_thresholds(model, X_test, y_test):
    """评估不同阈值下的性能"""
    print("\n=== 评估不同阈值下的性能 ===")
    
    y_proba = model.predict_proba(X_test)[:, 1]
    
    # 测试不同的阈值
    thresholds_to_test = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    
    results = []
    for threshold in thresholds_to_test:
        y_pred = (y_proba >= threshold).astype(int)
        
        # 计算混淆矩阵
        cm = confusion_matrix(y_test, y_pred)
        tn, fp, fn, tp = cm.ravel()
        
        # 计算各项指标
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred) if (tp + fp) > 0 else 0
        recall = recall_score(y_test, y_pred) if (tp + fn) > 0 else 0
        f1 = f1_score(y_test, y_pred) if (precision + recall) > 0 else 0
        specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
        
        results.append({
            '阈值': threshold,
            '准确率': accuracy,
            '精确率': precision,
            '召回率': recall,
            'F1得分': f1,
            '特异性': specificity,
            '真阳性': tp,
            '假阳性': fp,
            '真阴性': tn,
            '假阴性': fn
        })
    
    # 转换为DataFrame
    results_df = pd.DataFrame(results)
    print(results_df.round(4))
    
    return results_df

def plot_threshold_analysis(model, X_test, y_test):
    """绘制阈值分析图"""
    print("\n=== 生成阈值分析图 ===")
    
    y_proba = model.predict_proba(X_test)[:, 1]
    precision, recall, thresholds = precision_recall_curve(y_test, y_proba)
    
    # 创建子图
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # 1. Precision-Recall曲线
    axes[0, 0].plot(recall, precision, 'b-', linewidth=2)
    axes[0, 0].set_xlabel('召回率')
    axes[0, 0].set_ylabel('精确率')
    axes[0, 0].set_title('Precision-Recall曲线')
    axes[0, 0].grid(True, alpha=0.3)
    
    # 2. 阈值 vs 精确率和召回率
    axes[0, 1].plot(thresholds, precision[:-1], 'r-', label='精确率', linewidth=2)
    axes[0, 1].plot(thresholds, recall[:-1], 'b-', label='召回率', linewidth=2)
    axes[0, 1].set_xlabel('阈值')
    axes[0, 1].set_ylabel('指标值')
    axes[0, 1].set_title('阈值 vs 精确率/召回率')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)
    
    # 3. F1得分 vs 阈值
    f1_scores = 2 * (precision[:-1] * recall[:-1]) / (precision[:-1] + recall[:-1])
    f1_scores = np.nan_to_num(f1_scores)  # 处理除零错误
    axes[1, 0].plot(thresholds, f1_scores, 'g-', linewidth=2)
    axes[1, 0].set_xlabel('阈值')
    axes[1, 0].set_ylabel('F1得分')
    axes[1, 0].set_title('阈值 vs F1得分')
    axes[1, 0].grid(True, alpha=0.3)
    
    # 4. 预测概率分布
    axes[1, 1].hist(y_proba[y_test == 0], bins=20, alpha=0.7, label='正常', color='blue')
    axes[1, 1].hist(y_proba[y_test == 1], bins=20, alpha=0.7, label='异常', color='red')
    axes[1, 1].set_xlabel('预测概率')
    axes[1, 1].set_ylabel('频数')
    axes[1, 1].set_title('预测概率分布')
    axes[1, 1].legend()
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('threshold_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()

def apply_optimal_threshold_and_evaluate(model, X_train, X_test, y_train, y_test, threshold):
    """使用最优阈值进行预测和评估"""
    print(f"\n=== 使用最优阈值 {threshold:.4f} 进行评估 ===")
    
    # 训练集预测
    y_train_proba = model.predict_proba(X_train)[:, 1]
    y_train_pred = (y_train_proba >= threshold).astype(int)
    
    # 测试集预测
    y_test_proba = model.predict_proba(X_test)[:, 1]
    y_test_pred = (y_test_proba >= threshold).astype(int)
    
    # 计算各种指标
    results = {
        '数据集': ['训练集', '测试集'],
        '准确率': [
            accuracy_score(y_train, y_train_pred),
            accuracy_score(y_test, y_test_pred)
        ],
        '精确率': [
            precision_score(y_train, y_train_pred),
            precision_score(y_test, y_test_pred)
        ],
        '召回率': [
            recall_score(y_train, y_train_pred),
            recall_score(y_test, y_test_pred)
        ],
        'F1得分': [
            f1_score(y_train, y_train_pred),
            f1_score(y_test, y_test_pred)
        ],
        'AUC': [
            roc_auc_score(y_train, y_train_proba),
            roc_auc_score(y_test, y_test_proba)
        ]
    }
    
    results_df = pd.DataFrame(results)
    print("优化阈值后的性能指标:")
    print(results_df.round(4))
    
    # 打印测试集混淆矩阵
    cm_test = confusion_matrix(y_test, y_test_pred)
    print(f"\n测试集混淆矩阵:")
    print(f"真阴性: {cm_test[0,0]}, 假阳性: {cm_test[0,1]}")
    print(f"假阴性: {cm_test[1,0]}, 真阳性: {cm_test[1,1]}")
    
    # 计算特异性
    tn, fp, fn, tp = cm_test.ravel()
    specificity = tn / (tn + fp)
    print(f"\n额外指标:")
    print(f"特异性(正常样本正确识别率): {specificity:.4f}")
    print(f"敏感性(异常样本正确识别率): {tp/(tp+fn):.4f}")
    
    return results_df, y_test_pred, cm_test

def generate_threshold_optimization_report(results_df, cm_test, optimal_threshold, original_results):
    """生成阈值优化报告"""
    print("\n=== 生成阈值优化报告 ===")
    
    tn, fp, fn, tp = cm_test.ravel()
    specificity = tn / (tn + fp)
    sensitivity = tp / (tp + fn)
    
    report = f"""
# 分类阈值优化报告

## 阈值优化结果

### 最优阈值设置
- 优化后阈值: {optimal_threshold:.4f}
- 默认阈值: 0.5000

### 测试集性能对比

#### 优化后性能 (阈值={optimal_threshold:.4f})
- 准确率: {results_df.loc[1, '准确率']:.4f}
- 精确率: {results_df.loc[1, '精确率']:.4f}
- 召回率: {results_df.loc[1, '召回率']:.4f}
- F1得分: {results_df.loc[1, 'F1得分']:.4f}
- 特异性: {specificity:.4f}

#### 原始性能 (阈值=0.5000)
- 准确率: 0.8347
- 精确率: 0.8571  
- 召回率: 0.8000
- F1得分: 0.8276

### 性能提升分析
- 召回率提升: {results_df.loc[1, '召回率'] - 0.8000:.4f} ({((results_df.loc[1, '召回率'] - 0.8000) / 0.8000 * 100):+.1f}%)
- 精确率变化: {results_df.loc[1, '精确率'] - 0.8571:.4f} ({((results_df.loc[1, '精确率'] - 0.8571) / 0.8571 * 100):+.1f}%)
- 准确率变化: {results_df.loc[1, '准确率'] - 0.8347:.4f} ({((results_df.loc[1, '准确率'] - 0.8347) / 0.8347 * 100):+.1f}%)

### 混淆矩阵分析 (优化后)
- 真阴性 (TN): {tn} - 正确识别的正常样本
- 假阳性 (FP): {fp} - 误诊为异常的正常样本  
- 假阴性 (FN): {fn} - 漏诊的异常样本
- 真阳性 (TP): {tp} - 正确识别的异常样本

### 医学意义解释
- **敏感性 (召回率): {sensitivity:.4f}** - 在所有真实异常病例中，模型能正确识别的比例
- **特异性: {specificity:.4f}** - 在所有真实正常病例中，模型能正确识别的比例
- **漏诊率: {fn/(tp+fn):.4f}** - 异常病例被误诊为正常的比例  
- **误诊率: {fp/(tn+fp):.4f}** - 正常病例被误诊为异常的比例

## 结论和建议

### 优化效果
通过调整分类阈值到{optimal_threshold:.4f}，成功提高了模型的召回率，减少了漏诊风险。

### 实际应用建议
1. **医学筛查场景**: 提高召回率有助于减少漏诊，对高风险孕妇进行更全面的筛查
2. **资源配置**: 需要权衡误诊率上升带来的资源消耗
3. **临床决策**: 建议结合其他临床指标进行综合判断

### 风险提示
- 阈值降低会增加假阳性，可能导致不必要的进一步检查
- 建议在实际应用中进行更大样本的验证

生成时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    with open('threshold_optimization_report.txt', 'w', encoding='utf-8') as f:
        f.write(report)
    
    print("阈值优化报告已保存到: threshold_optimization_report.txt")
    return report

def main():
    """主函数"""
    print("=== 分类阈值优化脚本 ===")
    
    # 1. 加载数据
    X_train, X_test, y_train, y_test, feature_columns = load_processed_data()
    if X_train is None:
        return
    
    # 2. 训练优化模型
    model = train_optimized_random_forest(X_train, y_train)
    
    # 3. 评估不同阈值的性能
    threshold_results = evaluate_different_thresholds(model, X_test, y_test)
    
    # 4. 寻找最优阈值(目标召回率90%)
    optimal_threshold, optimal_recall, optimal_precision = find_optimal_threshold_for_recall(
        model, X_test, y_test, min_recall=0.90
    )
    
    # 如果无法达到90%召回率，尝试85%
    if optimal_recall < 0.90:
        print("\n尝试降低目标召回率到85%...")
        optimal_threshold, optimal_recall, optimal_precision = find_optimal_threshold_for_recall(
            model, X_test, y_test, min_recall=0.85
        )
    
    # 5. 使用最优阈值进行评估
    results_df, y_test_pred, cm_test = apply_optimal_threshold_and_evaluate(
        model, X_train, X_test, y_train, y_test, optimal_threshold
    )
    
    # 6. 生成可视化分析
    plot_threshold_analysis(model, X_test, y_test)
    
    # 7. 生成详细报告
    original_results = {'recall': 0.8000, 'precision': 0.8571, 'accuracy': 0.8347}
    report = generate_threshold_optimization_report(
        results_df, cm_test, optimal_threshold, original_results
    )
    
    # 8. 输出总结
    print("\n" + "="*60)
    print("阈值优化完成！")
    print(f"最优阈值: {optimal_threshold:.4f}")
    print(f"优化后召回率: {optimal_recall:.4f} ({optimal_recall:.1%})")
    print(f"优化后精确率: {optimal_precision:.4f} ({optimal_precision:.1%})")
    print("\n生成的文件:")
    print("📊 threshold_analysis.png - 阈值分析图")
    print("📄 threshold_optimization_report.txt - 阈值优化报告")
    print("="*60)

if __name__ == "__main__":
    main()